package com.zy;

import dev.langchain4j.data.image.Image;
import dev.langchain4j.model.image.ImageModel;
import dev.langchain4j.model.output.Response;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;

/**
 * @program: AI_langchain4j
 * @description: 自定义一个文生图模型， 联接的是硅基流动中的模型
 * @author: zy
 * @create: 2025-06-22 11:23
 */
public class _04_MyImageModelTest {
    public static void main(String[] args) {
        ImageModel imageModel=new SiliconFlowImageModel(   System.getenv("SILICON_API_KEY")    );
        Response<Image> response=imageModel.generate("一只很可爱的猫在草地上玩");
        System.out.println( "生成的图像地址:"+  response.content().url() );

    }
}

//自定义一个文生图模型
class SiliconFlowImageModel implements ImageModel {
    //这就是一个http客户端，用来发送http请求
    private final String apiKey;
    private final OkHttpClient client=new OkHttpClient();

    public SiliconFlowImageModel(String apiKey) {
        this.apiKey = apiKey;
    }

    @Override
    public Response<Image> generate(String prompt) {

        MediaType mediaType= MediaType.get("application/json; charset=utf-8");
        String json = String.format("{\"prompt\":\"%s\",\"model\":\"%s\", \"size\":\"512x512\"}", prompt, "Kwai-Kolors/Kolors");

        Request request=new Request.Builder()
                .url("https://api.siliconflow.cn/v1/images/generations")
                .post(RequestBody.create( json, mediaType   )  )
                .addHeader(   "Authorization","Bearer "+apiKey )
                .build();

       try( okhttp3.Response response= client.newCall( request ).execute(); ) {
           if (!response.isSuccessful()) {
               throw new RuntimeException("请求失败:" + response.code());
           }
           String responseBody = response.body().string();

           //取数据    // 假设返回的是 {"image_url": "http..."}
           int start = responseBody.indexOf("http");
           int end = responseBody.indexOf("\"", start);
           String url = responseBody.substring(start, end);

           Image image = Image.builder().url(url).build();
           Response<Image> result = Response.from(image);

           return result;
       }catch(Exception ex){
           ex.printStackTrace();;
           throw new RuntimeException("图像生成异常:"+ ex.getMessage() );
       }
    }
}
